#include "custom_hough_lines.h"
#include <opencv2/imgproc/imgproc.hpp>

namespace rmd {

CustomHoughLineP::CustomHoughLineP() {}

CustomHoughLineP::~CustomHoughLineP() {}

void CustomHoughLineP::buuHoughLinesProbabilistic(CvMat* image, int threshold,
                                                  int lineLength, int lineGap,
                                                  std::vector<cv::Vec4i>& lines,
                                                  int linesMax, int startY,
                                                  int endY) {
  float rho = 1.f;
  int mtheta = 180;

  cv::Mat accum, mask;
  cv::MemStorage storage(cvCreateMemStorage(0));
  CvSeq* seq;
  CvSeqWriter writer;
  int width, height;
  int num_angle, num_rho;
  int r, count;
  CvPoint pt;
  CvRNG rng = cvRNG(-1);
  const float* ttab;
  uchar* mdata0;
  CV_Assert(CV_IS_MAT(image) && CV_MAT_TYPE(image->type) == CV_8UC1);

  width = image->cols;
  height = image->rows;

  startY = startY < 0 ? 0 : startY;
  startY = startY > height ? height : startY;

  endY = endY < 0 ? 0 : endY;
  endY = endY > height ? height : endY;

  if (lineGap > endY - startY) {
    lineGap = endY - startY;
  }

  // num_angle = cvRound(CV_PI / theta);
  num_angle = mtheta;
  num_rho = cvRound((((width + height) << 1) + 1) / rho);

  accum.create(num_angle, num_rho, CV_32SC1);
  mask.create(height, width, CV_8UC1);
  std::vector<float> trigtab(num_angle << 1);
  // trigtab.resize(num_angle*2);
  accum = cv::Scalar(0);
  for (int n = 0; n < num_angle; ++n) {
    trigtab[n << 1] = cosValue[n];
    trigtab[(n << 1) + 1] = sinValue[n];
  }

  /*
  for (ang = 0, n = 0; n < num_angle; ang += theta, n++)
  {
      trigtab[n * 2] = (float)(cos(ang) * irho);
      trigtab[n * 2 + 1] = (float)(sin(ang) * irho);
  }*/

  ttab = &trigtab[0];
  mdata0 = mask.data;
  cvStartWriteSeq(CV_32SC2, sizeof(CvSeq), sizeof(CvPoint), storage, &writer);

  // for (pt.y = startY; pt.y < endY; ++pt.y)
  for (pt.y = 0; pt.y < height; ++pt.y) {
    const uchar* data = image->data.ptr + pt.y * image->step;
    uchar* mdata = mdata0 + pt.y * width;
    for (pt.x = 0; pt.x < width; ++pt.x) {
      if (data[pt.x]) {
        mdata[pt.x] = (uchar)1;
        CV_WRITE_SEQ_ELEM(pt, writer);
      } else {
        mdata[pt.x] = 0;
      }
    }
  }

  seq = cvEndWriteSeq(&writer);
  count = seq->total;

  for (; count > 0; count--) {
    int idx = cvRandInt(&rng) % count;
    int max_val = threshold - 1;
    int max_n = 0;
    CvPoint* point = (CvPoint*)cvGetSeqElem(seq, idx);
    CvPoint line_end[2] = {(0, 0), (0, 0)};
    float a, b;
    int* adata = (int*)accum.data;
    int i, j, x0, y0, dx0, dy0, xflag;
    int good_line;
    const int shift = 16;
    i = point->y;
    j = point->x;

    *point = *(CvPoint*)cvGetSeqElem(seq, count - 1);

    if (!mdata0[i * width + j]) {
      continue;
    }

    for (int n = 0; n < num_angle; n++, adata += num_rho) {
      r = cvRound(j * ttab[n << 1] + i * ttab[(n << 1) + 1]);
      r += (num_rho - 1) >> 1;
      int val = ++adata[r];
      if (max_val < val) {
        max_val = val;
        max_n = n;
      }
    }

    if (max_val < threshold) {
      continue;
    }

    a = -ttab[(max_n << 1) + 1];  // a=-sin
    b = ttab[max_n << 1];         // b=cos
    x0 = j;
    y0 = i;

    if (fabs(a) > fabs(b)) {
      xflag = 1;
      dx0 = a > 0 ? 1 : -1;
      dy0 = cvRound(b * (1 << shift) / fabs(a));
      y0 = (y0 << shift) + (1 << (shift - 1));
    } else {
      xflag = 0;
      dy0 = b > 0 ? 1 : -1;
      dx0 = cvRound(a * (1 << shift) / fabs(b));
      x0 = (x0 << shift) + (1 << (shift - 1));
    }

    for (int k = 0; k < 2; ++k) {
      int gap = 0, x = x0, y = y0, dx = dx0, dy = dy0;

      if (k > 0) {
        dx = -dx;
        dy = -dy;
      }

      for (;; x += dx, y += dy) {
        uchar* mdata;
        int i1, j1;
        if (xflag) {
          j1 = x;
          i1 = y >> shift;
        } else {
          j1 = x >> shift;
          i1 = y;
        }
        if (j1 < 0 || j1 >= width || i1 < 0 || i1 >= height) {
          break;
        }
        mdata = mdata0 + i1 * width + j1;

        if (*mdata) {
          gap = 0;
          line_end[k].y = i1;
          line_end[k].x = j1;
        } else if (++gap > lineGap) {
          break;
        }
      }
    }
    good_line = (abs(line_end[1].x - line_end[0].x) >= lineLength) ||
                (abs(line_end[1].y - line_end[0].y) >= lineLength);

    //    std::vector<cv::Point> temp_point_serial;

    for (int k = 0; k < 2; ++k) {
      int x = x0, y = y0, dx = dx0, dy = dy0;
      if (k > 0) {
        dx = -dx;
        dy = -dy;
      }

      for (;; x += dx, y += dy) {
        uchar* mdata;
        int i1, j1;

        if (xflag) {
          j1 = x;
          i1 = y >> shift;
        } else {
          j1 = x >> shift;
          i1 = y;
        }

        mdata = mdata0 + i1 * width + j1;

        if (*mdata) {
          if (good_line) {
            CvPoint good_line_point;
            good_line_point.y = i1;
            good_line_point.x = j1;
            //            temp_point_serial.push_back(good_line_point);

            adata = (int*)accum.data;
            for (int n = 0; n < num_angle; n++, adata += num_rho) {
              r = cvRound(j1 * ttab[n << 1] + i1 * ttab[(n << 1) + 1]);
              r += (num_rho - 1) >> 1;
              adata[r]--;
            }
          }
          *mdata = 0;
        }
        if (i1 == line_end[k].y && j1 == line_end[k].x) {
          break;
        }
      }
    }
    //    int temp_point_size = temp_point_serial.size();
    if (good_line) {
      lines.push_back(cv::Vec4i(line_end[0].x, line_end[0].y, line_end[1].x,
                                line_end[1].y));
      //      vec2LinePoint.push_back(temp_point_serial);

      if (lines.size() >= linesMax) {
        return;
      }
    }
  }
}

void CustomHoughLineP::buuHoughLinesP(cv::Mat& image,
                                      std::vector<cv::Vec4i>& _lines,
                                      int threshold, double minLineLength,
                                      double maxGap, int startY, int endY) {
  CvMat c_image = image;
  buuHoughLinesProbabilistic(&c_image, threshold, cvRound(minLineLength),
                             cvRound(maxGap), _lines, SHRT_MAX, startY, endY);
}
}
